package com.qny.ai.websocket;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentSkipListSet;

@Component
public class SignalingHandler extends TextWebSocketHandler {

    private final ObjectMapper mapper = new ObjectMapper();
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>(); // sessionId -> session
    private final Map<String, String> sessionToUser = new ConcurrentHashMap<>();      // sessionId -> userId
    private final Map<String, WebSocketSession> userToSession = new ConcurrentHashMap<>(); // userId -> session
    private final Map<String, Set<String>> roomToUsers = new ConcurrentHashMap<>();   // roomId -> userIds

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        sessions.put(session.getId(), session);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        JsonNode node = mapper.readTree(message.getPayload());
        String type = node.path("type").asText();
        String roomId = node.path("roomId").asText("");
        String fromUserId = node.path("from").asText("");
        String toUserId = node.path("to").asText("");

        switch (type) {
            case "join" -> {
                if (!fromUserId.isEmpty()) {
                    sessionToUser.put(session.getId(), fromUserId);
                    userToSession.put(fromUserId, session);
                }
                if (!roomId.isEmpty()) {
                    roomToUsers.computeIfAbsent(roomId, k -> new ConcurrentSkipListSet<>()).add(fromUserId);
                }
                // 回发确认
                session.sendMessage(new TextMessage("{\"type\":\"joined\",\"roomId\":\"" + roomId + "\"}"));
            }
            case "leave" -> {
                if (!roomId.isEmpty() && !fromUserId.isEmpty()) {
                    Set<String> set = roomToUsers.get(roomId);
                    if (set != null) set.remove(fromUserId);
                }
                session.sendMessage(new TextMessage("{\"type\":\"left\",\"roomId\":\"" + roomId + "\"}"));
            }
            case "offer", "answer", "candidate" -> {
                // 定向转发给指定用户
                if (!toUserId.isEmpty()) {
                    WebSocketSession target = userToSession.get(toUserId);
                    if (target != null && target.isOpen()) {
                        target.sendMessage(new TextMessage(message.getPayload()));
                    } else {
                        session.sendMessage(new TextMessage("{\"type\":\"error\",\"message\":\"peer not found\"}"));
                    }
                } else if (!roomId.isEmpty()) {
                    // 房间内广播（排除自己）
                    Set<String> users = roomToUsers.get(roomId);
                    if (users != null) {
                        for (String uid : users) {
                            if (!uid.equals(fromUserId)) {
                                WebSocketSession target = userToSession.get(uid);
                                if (target != null && target.isOpen()) {
                                    target.sendMessage(new TextMessage(message.getPayload()));
                                }
                            }
                        }
                    }
                }
            }
            default -> session.sendMessage(new TextMessage("{\"type\":\"error\",\"message\":\"unknown type\"}"));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        sessions.remove(session.getId());
        String uid = sessionToUser.remove(session.getId());
        if (uid != null) {
            userToSession.remove(uid);
            // 从所有房间移除该用户
            for (Set<String> set : roomToUsers.values()) {
                set.remove(uid);
            }
        }
    }
}


